import torch
import torch.nn as nn


class VAEModel(nn.Module):
    def __init__(self, x_dim, h_dim1, h_dim2, z_dim):
        super().__init__()
        self.x_dim = x_dim
        # 编码器
        self.fc1 = nn.Linear(x_dim, h_dim1)
        self.fc2 = nn.Linear(h_dim1, h_dim2)  # 全连接层
        self.fc31 = nn.Linear(h_dim2, z_dim)  # 全连接层
        self.fc32 = nn.Linear(h_dim2, z_dim)

        # 解码器
        self.fc4 = nn.Linear(z_dim, h_dim2)
        self.fc5 = nn.Linear(h_dim2, h_dim1)
        self.fc6 = nn.Linear(h_dim1, x_dim)

    def encoder(self, x):
        h = torch.relu(self.fc1(x))
        h = torch.relu(self.fc2(h))
        return self.fc31(h), self.fc32(h)

    def decoder(self, z):
        h = torch.relu(self.fc4(z))
        h = torch.relu(self.fc5(h))
        return torch.sigmoid(self.fc6(h))

    # 重参数化技巧
    def sampling(self, mu, log_var):
        # 计算标准淮差
        std = torch.exp(0.5 * log_var)  # 从标准正态分布中随机采样eps
        eps = torch.randn_like(std)  # 返回z
        return mu + eps * std

    def forward(self, x):
        mu, log_var = self.encoder(x.view(-1, self.x_dim))
        z = self.sampling(mu, log_var)
        return self.decoder(z), mu, log_var, z


if __name__ == '__main__':
    vae = VAEModel(28 * 28, 512, 256, 2)
